from fastapi import FastAPI, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from database import Base, engine, DATABASE_URL
from auth import router as auth_router
from stock import get_stock_router
from market import get_market_router
from utils import load_parquet
import os
from contextlib import asynccontextmanager
import akshare as ak
import pandas as pd
from retrying import retry
from router import router

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifecycle"""
    # Initialize user database
    db_path = DATABASE_URL.replace("sqlite+aiosqlite:///", "")
    if not os.path.exists(db_path):
        print("⚙️ 用户数据库文件不存在，正在初始化...")
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        print("✅ 用户数据库初始化完成。")
    else:
        print("🗃 用户数据库已存在，跳过初始化。")

    # Load parquet data
    try:
        app.state.df = load_parquet()
        print("Parquet 数据加载成功")
    except FileNotFoundError as e:
        print(f"错误: {e}")
        app.state.df = None
    except Exception as e:
        print(f"加载 Parquet 文件时发生错误: {e}")
        app.state.df = None

    yield
    # Cleanup
    app.state.df = None

app = FastAPI(
    title="股票数据与用户管理API服务",
    lifespan=lifespan
)

# CORS Middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Retry decorator for AKShare
@retry(stop_max_attempt_number=3, wait_fixed=2000)
def fetch_index_data(start_date_ak, end_date_ak):
    """Fetches index data from AKShare with retry logic."""
    return ak.index_zh_a_hist(symbol="932000", start_date=start_date_ak, end_date=end_date_ak)

# Include existing Routers
app.include_router(auth_router)
stock_router = get_stock_router(app)  # Updated to only get router
app.include_router(stock_router)
app.include_router(router)
app.include_router(get_market_router(app))

if __name__ == "__main__":
    import uvicorn
    os.makedirs('job/data', exist_ok=True)
    os.makedirs('data/backtest_results', exist_ok=True)
    uvicorn.run(app, host="0.0.0.0", port=8099)